import time

import matplotlib
import numpy as np
from matplotlib import pyplot as plt

from utils import save_variable, analyze

matplotlib.use("TkAgg")

fig = plt.figure()
ax = fig.subplots()


def F(x):
    """
    目标函数

    :param x: 输入
    :return: 函数输出
    """
    return x * np.sin(10 * np.pi * x) + 1.0


def move(x) -> float:
    """
    通过正态分布随机移动一步x

    :param x: x
    :return: 移动后的x
    """
    rand = np.random.randn() / 1.0
    return x + rand


class SA:
    def __init__(self, function, f_range, init_t, min_t, alpha_t, k):
        """
        模拟退火算法参数初始化

        :param function: 目标函数
        :param f_range: 函数范围
        :param init_t: 初始温度
        :param min_t: 终止温度
        :param alpha_t: 温度衰减系数
        :param k: 温度衰减间隔迭代次数
        """
        self.init_t = init_t
        self.min_t = min_t
        self.k = k
        self.F = function
        self.range_l = f_range[0]
        self.range_r = f_range[1]
        self.alpha_t = alpha_t
        self.plot = True
        self.show_log = True

    def search(self):
        """
        执行模拟退火算法

        :return: 函数最大值处的x
        """
        t = self.init_t
        x = np.random.random_sample() * (self.range_r - self.range_l) + self.range_l
        if self.show_log:
            print("x_init =", x)
        x_f = np.linspace(-1, 2, num=300)
        y_f = self.F(x_f)
        iteration = 1

        # t小于终止温度，结束搜索
        while t > self.min_t:
            avg_p = 0.0
            count_p = 0
            # 每k次迭代降低温度
            for i in range(self.k):
                x_new = move(x)
                if x_new < self.range_l or x_new > self.range_r:
                    continue
                y = self.F(x)
                y_new = self.F(x_new)
                # 比当前值小，直接采纳
                if y_new < y:
                    x = x_new
                # 比当前值大，一定概率采纳
                else:
                    p = np.exp(-(y_new - y) / t)
                    avg_p += p
                    count_p += 1
                    if p > np.random.random_sample():
                        x = x_new

                # 绘制动画
                if self.plot:
                    ax.cla()
                    ax.plot(x_f, y_f)
                    ax.axvline(x, color='r')
                    ax.plot(x, self.F(x), marker="o", markersize=5, markeredgecolor="green")
                    plt.pause(0.0001)

            # 计算该温度下的平均概率
            avg_p = avg_p / count_p

            if self.show_log:
                print("epoch", iteration * self.k, ": t =", t, ", avg_p =", avg_p)
            t = t * self.alpha_t
            iteration += 1

        return x


if __name__ == "__main__":
    sa = SA(lambda x: -F(x), f_range=[-1.0, 2.0], init_t=1, min_t=0.015, alpha_t=0.93, k=50)

    # 模拟一次
    result = sa.search()
    print("result:", "x_max =", result, ", y_max =", F(result))

    # 模拟多次
    # sa.show_log = False
    # sa.plot = False
    # x_list = []
    # y_list = []
    # count = 5000
    # start = time.time()
    # for i in range(count):
    #     x_max = sa.search()
    #     x_list.append(x_max)
    #     y_list.append(F(x_max))
    # end = time.time()
    # save_variable(x_list, "x_list_SA.pkl")
    # save_variable(y_list, "y_list_SA.pkl")
    # info = {
    #     "consume": end - start,
    #     "best_x": max(x_list, key=F),
    #     "best_y": max(y_list),
    #     "worst_y": min(y_list)
    # }
    # save_variable(info, "info_SA.pkl")
